#ifndef SRC_HELIB_ATM_WRAPPER
#define SRC_HELIB_ATM_WRAPPER
/* Copyright (C) 2012-2020 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
#include <helib/fhe_stats.h>
#include <helib/helib.h>

namespace helib {
/**
 * @class BasicAutomorphPrecon
 * @brief Pre-computation to speed many automorphism on the same ciphertext.
 *
 * The expensive part of homomorphic automorphism is breaking the ciphertext
 * parts into digits. The usual setting is we first rotate the ciphertext
 * parts, then break them into digits. But when we apply many automorphisms
 * it is faster to break the original ciphertext into digits, then rotate
 * the digits (as opposed to first rotate, then break).
 * An BasicAutomorphPrecon object breaks the original ciphertext and keeps
 * the digits, then when you call automorph is only needs to apply the
 * native automorphism and key switching to the digits, which is fast(er).
 **/
class BasicAutomorphPrecon {
  Ctxt ctxt;
  NTL::xdouble noise;
  std::vector<DoubleCRT> polyDigits;

public:
  BasicAutomorphPrecon(const Ctxt &_ctxt) : ctxt(_ctxt), noise(1.0) {
    HELIB_TIMER_START;
    if (ctxt.parts.size() >= 1)
      assertTrue(
          ctxt.parts[0].skHandle.isOne(),
          "Invalid ciphertext (secret key handle for part 0 is not one)");
    if (ctxt.parts.size() <= 1)
      return; // nothing to do

    ctxt.cleanUp();
    const Context &context = ctxt.getContext();
    const PubKey &pubKey = ctxt.getPubKey();
    long keyID = ctxt.getKeyID();

    // The call to cleanUp() should ensure that this assertion passes.
    assertTrue(ctxt.inCanonicalForm(keyID),
               "Ciphertext is not in canonical form");

    ctxt.relin_CKKS_adjust();

    // Compute the number of digits that we need and the estimated
    // added noise from switching this ciphertext.

    NTL::xdouble addedNoise = ctxt.parts[1].breakIntoDigits(polyDigits);
    NTL::xdouble max_ks_noise(0.0);
    for (const KeySwitch &ks : pubKey.keySWlist()) {
      if (max_ks_noise < ks.noiseBound)
        max_ks_noise = ks.noiseBound;
    }
    addedNoise *= max_ks_noise;

    double logProd = context.logOfProduct(context.getSpecialPrimes());
    noise = ctxt.getNoiseBound() * NTL::xexp(logProd);

    double ratio = NTL::conv<double>(addedNoise / noise);

    HELIB_STATS_UPDATE("KS-noise-ratio-hoist", ratio);
    // if (ratio > 1) {
    //   Warning("KS-noise-ratio-hoist=" + std::to_string(ratio));
    // }
    // std::stderr << "*** HOIST INIT\n";
    // fprintf(stderr, "   KS-log-noise-ratio-hoist: %f\n",
    // log(addedNoise/noise)/log(2.0));

    noise += addedNoise;
  }

  std::shared_ptr<Ctxt> automorph(long k) const {
    HELIB_TIMER_START;

    // A hack: record this automorphism rather than actually performing it
    if (isSetAutomorphVals()) { // defined in NumbTh.h
      recordAutomorphVal(k);
      return std::make_shared<Ctxt>(ctxt);
    }

    if (k == 1 || ctxt.isEmpty())
      return std::make_shared<Ctxt>(ctxt); // nothing to do

    const Context &context = ctxt.getContext();
    const PubKey &pubKey = ctxt.getPubKey();

    // empty ctxt
    std::shared_ptr<Ctxt> result = std::make_shared<Ctxt>(ZeroCtxtLike, ctxt);
    result->noiseBound = noise; // noise estimate
    result->intFactor = ctxt.intFactor;

    result->primeSet = ctxt.primeSet | context.getSpecialPrimes();
    // VJS-NOTE: added this to make addPart work

    if (ctxt.isCKKS()) {
      result->ptxtMag = ctxt.ptxtMag;
      double logProd = context.logOfProduct(context.getSpecialPrimes());
      result->ratFactor = ctxt.ratFactor * NTL::xexp(logProd);
    }

    if (ctxt.parts.size() == 1) { // only constant part, no need to key-switch
      CtxtPart tmpPart = ctxt.parts[0];
      tmpPart.automorph(k);
      tmpPart.addPrimesAndScale(context.getSpecialPrimes());
      result->addPart(tmpPart, /*matchPrimeSet=*/true);
      return result;
    }

    // Ensure that we have a key-switching matrices for this automorphism
    long keyID = ctxt.getKeyID();
    if (!pubKey.isReachable(k, keyID)) {
      throw LogicError("no key-switching matrices for k=" + std::to_string(k) +
                       ", keyID=" + std::to_string(keyID));
    }

    // Get the first key-switching matrix for this automorphism
    const KeySwitch &W = pubKey.getNextKSWmatrix(k, keyID);
    long amt = W.fromKey.getPowerOfX();

    // Start by rotating the constant part, no need to key-switch it
    CtxtPart tmpPart = ctxt.parts[0];
    tmpPart.automorph(amt);
    tmpPart.addPrimesAndScale(context.getSpecialPrimes());
    result->addPart(tmpPart, /*matchPrimeSet=*/true);

    // Then rotate the digits and key-switch them
    std::vector<DoubleCRT> tmpDigits = polyDigits;
    for (auto &&tmp : tmpDigits) // rotate each of the digits
      tmp.automorph(amt);

    result->keySwitchDigits(W, tmpDigits); // key-switch the digits

    long m = context.getM();
    if ((amt - k) % m != 0) { // amt != k (mod m), more automorphisms to do
      k = NTL::MulMod(k, NTL::InvMod(amt, m), m); // k *= amt^{-1} mod m
      result->smartAutomorph(k);                  // call usual smartAutomorph
    }
    return result;
  }
};

class GeneralAutomorphPrecon {
public:
  virtual ~GeneralAutomorphPrecon() {}

  virtual std::shared_ptr<Ctxt> automorph(long i) const = 0;
};

class GeneralAutomorphPrecon_FULL : public GeneralAutomorphPrecon {
private:
  BasicAutomorphPrecon precon;
  long dim;
  const PAlgebra &zMStar;

public:
  GeneralAutomorphPrecon_FULL(const Ctxt &_ctxt, long _dim,
                              const EncryptedArray &ea)
      : precon(_ctxt), dim(_dim), zMStar(ea.getPAlgebra()) {}

  std::shared_ptr<Ctxt> automorph(long i) const override {
    return precon.automorph(zMStar.genToPow(dim, i));
  }
};

} // namespace helib
#endif /* SRC_HELIB_ATM_WRAPPER */
